#include <cstdio>
#include <systemc>
#include <vector>
#include <random>
#include <map>
#include <cassert>
#include <cstdint>
#include <string>

using namespace sc_core;
using namespace sc_dt;
using namespace std;

// Define Master Interface (Bus provides this to Master)
class MasterIF : virtual public sc_interface {
public:
    virtual void write(uint8_t addr, uint32_t data, int id) = 0;
    virtual void read(uint8_t addr, uint32_t &data, int id) = 0;
};

// Define Slave Interface (Slave provides this to Bus)
class SlaveIF : virtual public sc_interface {
public:
    virtual void write(uint8_t addr, uint32_t data) = 0;
    virtual void read(uint8_t addr, uint32_t &data) = 0;
    virtual void get_map(size_t& start, size_t& size) = 0;
};

// Define Master Module
class Master : public sc_module {
private:
    int id;
public:
    // Master -> Bus
    sc_port<MasterIF> bus_port;

    // Constructor
    Master(sc_module_name _name, int _id) : sc_module(_name), id(_id) { }

    // Write operation (called from Top module)
    void write(uint8_t addr, uint32_t data) {
        printf("@%s [%s]:\tWRITE REQUEST  addr=0x%02x data=0x%08x\n", sc_time_stamp().to_string().c_str(), this->name(), addr, data);
        bus_port->write(addr, data, id);
        printf("@%s [%s]:\tWRITE RESPONSE addr=0x%02x data=0x%08x\n", sc_time_stamp().to_string().c_str(), this->name(), addr, data);
    }

    // Read operation (called from Top module)
    void read(uint8_t addr, uint32_t& data) {
        printf("@%s [%s]:\tREAD  REQUEST  addr=0x%02x\n", sc_time_stamp().to_string().c_str(), this->name(), addr);
        bus_port->read(addr, data, id);
        printf("@%s [%s]:\tREAD  RESPONSE addr=0x%02x data=0x%08x\n", sc_time_stamp().to_string().c_str(), this->name(), addr, data);
    }
};

// Define Bus Module
template <int N_MASTERS, int N_SLAVES>
class Bus : public sc_module, public MasterIF {
private:
    // For arbitration
    bool request[N_MASTERS] = {false};
    sc_event proceed[N_MASTERS];
    int last_granted = -1;
public:
    // Clock signal
    sc_in<bool> clock;
    // Bus -> Slaves
    sc_port<SlaveIF, N_SLAVES> slave_ports;
    // Masters <- Bus
    sc_export<MasterIF> master_export;

    // Constructor
    Bus(sc_module_name _name) : sc_module(_name) {
        master_export.bind(*this);
        SC_METHOD(arbitrate);
        sensitive << clock.pos();
        dont_initialize();
    }

    // Find slave port by address
    int find_port(uint8_t addr) {
        for (int i = 0; i < N_SLAVES; ++i) {
            size_t start, size;
            slave_ports[i]->get_map(start, size);
            if (addr >= start && addr < start + size) {
                return i;
            }
        }
        SC_REPORT_ERROR("Bus", ("Address " + to_string(addr) + " not mapped to any slave").c_str());
        return -1;
    }

    // Write operation (implementing master_if, called by Master)
    void write(uint8_t addr, uint32_t data, int id) override {
        request[id] = true;
        wait(proceed[id]); // Wait for arbitration grant
        request[id] = false;

        int port_idx = find_port(addr);
        if (port_idx != -1) {
            slave_ports[port_idx]->write(addr, data);
        }
    }

    // Read operation (implementing master_if, called by Master)
    void read(uint8_t addr, uint32_t &data, int id) override {
        request[id] = true;
        wait(proceed[id]); // Wait for arbitration grant
        request[id] = false;

        int port_idx = find_port(addr);
        if (port_idx != -1) {
            slave_ports[port_idx]->read(addr, data);
        }
    }

    // Round-robin arbitration
    void arbitrate() {
        int start = (last_granted + 1) % N_MASTERS;
        for (int i = 0; i < N_MASTERS; ++i) {
            int idx = (start + i) % N_MASTERS;
            if (request[idx]) {
                proceed[idx].notify(); // Notify the waiting master thread
                last_granted = idx;
                break;
            }
        }
    }
};

// Define Slave Module
class Slave : public sc_module, public SlaveIF {
private:
    vector<uint32_t> memory;
    size_t saddr;
    size_t msize;
public:
    // Bus <- Slave
    sc_export<SlaveIF> slave_export;

    // Constructor
    Slave(sc_module_name _name, size_t _saddr, size_t _msize)
        : sc_module(_name), saddr(_saddr), msize(_msize) {
        slave_export.bind(*this);
        memory.resize(msize);
    }

    // Write operation (implementing slave_if, called by Bus)
    void write(uint8_t addr, uint32_t data) override {
        assert(addr >= saddr && addr - saddr < msize);
        memory[addr - saddr] = data;
        printf("@%s [%s]:\tWRITE EVENT    addr=0x%02x data=0x%08x\n", sc_time_stamp().to_string().c_str(), this->name(), addr, data);
    }

    // Read operation (implementing slave_if, called by Bus)
    void read(uint8_t addr, uint32_t &data) override {
        assert(addr >= saddr && addr - saddr < msize);
        data = memory[addr - saddr];
        printf("@%s [%s]:\tREAD  EVENT    addr=0x%02x data=0x%08x\n", sc_time_stamp().to_string().c_str(), this->name(), addr, data);
    }

    // Get memory map (implementing slave_if, called by Bus)
    void get_map(size_t& start, size_t& size) override {
        start = saddr;
        size = msize;
    }
};


// Define Top Module
template <size_t N_MASTERS, size_t N_SLAVES>
class Top: public sc_module {
private:
    // Components
    Bus<N_MASTERS, N_SLAVES>* bus;
    Master* masters[N_MASTERS];
    Slave* slaves[N_SLAVES];
    // Record of write operations for verification
    map<uint8_t, uint32_t> write_record;
    // Random number generator
    std::mt19937 gen;
    std::uniform_int_distribution<uint32_t> data_dist;
    std::uniform_int_distribution<uint8_t> addr_dist;
    std::uniform_int_distribution<uint8_t> rwop_dist; // 0: read, 1: write
public:
    // Clock signal
    sc_in<bool> clock;

    // Constructor
    Top(sc_module_name _name):
        sc_module(_name), gen(0), data_dist(0, 0xFFFFFFFF), addr_dist(0, 255), rwop_dist(0, 1) {

        // Initialize bus
        bus = new Bus<N_MASTERS, N_SLAVES>("bus");
        bus->clock(clock);

        // Initialize masters
        for (int i = 0; i < N_MASTERS; ++i) {
            masters[i] = new Master(("master" + to_string(i)).c_str(), i);
            masters[i]->bus_port.bind(bus->master_export);
        }

        // Initialize slaves
        unsigned int addr_step = 256 / N_SLAVES;
        for (int i = 0; i < N_SLAVES; ++i) {
            slaves[i] = new Slave(("slave" + to_string(i)).c_str(), i * addr_step, addr_step);
            bus->slave_ports.bind(slaves[i]->slave_export);
        }

        // Spawn processes for each master
        for (int i = 0; i < N_MASTERS; ++i) {
            sc_spawn([this, i]() {
                this->run(i);
            });
        }
        printf("Top constructor called\n");
    }

    void run(int id) {
        uint8_t addr;
        uint32_t data;

        while (true) {
            // Send request at every negative edge of the clock
            wait(clock.negedge_event());
            uint8_t op = rwop_dist(gen);

            // Force write if no record
            if (write_record.empty()) {
                op = 1;
            }

            // Read operation
            if (op == 0) {
                // Randomly pick an address from recorded writes
                assert(!write_record.empty());
                auto it = write_record.begin();
                std::advance(it, gen() % write_record.size());
                addr = it->first;

                // Perform read
                data = 0;
                masters[id]->read(addr, data);

                // Verify data
                uint32_t expected_data = it->second; // Must get expected data after read !!!
                if (data != expected_data) {
                    printf("@%s [%s]: =====>\t VERIFY FAILED at addr=0x%02x: expected=0x%08x, got=0x%08x\n", sc_time_stamp().to_string().c_str(), this->name(), addr, expected_data, data);
                    assert(false);
                } else {
                    printf("@%s [%s]: =====>\tVERIFY PASSED at addr=0x%02x: data=0x%08x\n", sc_time_stamp().to_string().c_str(), this->name(), addr, data);
                }
            }
            // Write operation
            else if (op == 1) {
                // Randomly generate address and data
                addr = addr_dist(gen);
                data = data_dist(gen);
                // Perform write
                masters[id]->write(addr, data);
                // Record the write operation
                write_record[addr] = data;
            }
        }
    }
    ~Top() {
        delete bus;
        for (int i = 0; i < N_MASTERS; ++i) {
            delete masters[i];
        }
        for (int i = 0; i < N_SLAVES; ++i) {
            delete slaves[i];
        }
        printf("Top destructor called\n");
    }
};

int sc_main(int argc, char **argv) {

    const int N_MASTERS = 2;
    const int N_SLAVES = 2;
    const int period = 2;

    // Clock signal
    sc_clock clock("clock", period, SC_NS, 0.5, period>>1, SC_NS, true);

    // Instantiate Top Module
    Top<N_MASTERS, N_SLAVES> top("top");
    top.clock(clock);

    // Start simulation
    sc_start(500000 * period, SC_NS);

    return 0;
}
